#!/usr/bin/env python3
"""
Prometheus API数据访问模块
支持告警查询和状态分析
"""

import requests
import logging
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional
from config import Config

logger = logging.getLogger(__name__)

class PrometheusConnector:
    """Prometheus API连接器"""
    
    def __init__(self):
        self.config = Config.get_prometheus_config()
        self.base_url = self.config['prometheus_url']
        self.timeout = self.config['query_timeout']
        self.max_alerts = self.config['max_alerts']
        
        logger.info(f"初始化Prometheus连接器: {self.base_url}")
    
    def test_connection(self) -> bool:
        """测试Prometheus连接"""
        try:
            response = requests.get(
                f"{self.base_url}/api/v1/query",
                params={'query': 'up'},
                timeout=10
            )
            
            if response.status_code == 200:
                logger.info("Prometheus连接测试成功")
                return True
            else:
                logger.error(f"Prometheus连接失败: HTTP {response.status_code}")
                return False
                
        except Exception as e:
            logger.error(f"Prometheus连接异常: {e}")
            return False
    
    def get_active_alerts(self) -> List[Dict[str, Any]]:
        """获取当前活跃告警"""
        try:
            # 查询AlertManager的告警
            alerts_url = f"{self.base_url.rstrip('/api/v1')}/api/v1/alerts"
            
            logger.debug(f"请求告警URL: {alerts_url}")
            response = requests.get(alerts_url, timeout=self.timeout)
            
            if response.status_code == 200:
                data = response.json()
                logger.debug(f"原始API返回数据类型: {type(data)}, 内容: {str(data)[:200]}...")
                
                # Prometheus API返回的数据结构是 data.alerts，不是直接的 data
                if 'data' in data and 'alerts' in data['data']:
                    alerts = data['data']['alerts']
                elif 'data' in data and isinstance(data['data'], list):
                    # 如果 data 直接是数组
                    alerts = data['data']
                else:
                    logger.warning(f"未识别的API返回数据结构: {data}")
                    alerts = []
                
                logger.debug(f"解析得到告警数量: {len(alerts)}")
                
                # 过滤活跃告警并标准化数据
                active_alerts = []
                for i, alert in enumerate(alerts):
                    logger.debug(f"处理第{i+1}个告警: 数据类型={type(alert)}, state={alert.get('state') if isinstance(alert, dict) else 'N/A'}")
                    
                    if isinstance(alert, dict) and alert.get('state') == 'firing':
                        processed_alert = self._process_alert_data(alert)
                        active_alerts.append(processed_alert)
                        logger.debug(f"已处理告警: {processed_alert.get('alert_name', 'unknown')}")
                
                logger.info(f"获取到 {len(active_alerts)} 个活跃告警")
                return active_alerts[:self.max_alerts]
            else:
                logger.error(f"获取告警失败: HTTP {response.status_code}, 响应: {response.text[:200]}")
                return []
                
        except Exception as e:
            logger.error(f"获取告警异常: {e}")
            import traceback
            logger.error(f"异常堆栈: {traceback.format_exc()}")
            return []
    
    def get_metrics_summary(self) -> Dict[str, Any]:
        """获取系统指标摘要"""
        try:
            queries = {
                'total_targets': 'count(up)',
                'down_targets': 'count(up == 0)',
                'cpu_usage': 'avg(100 - (avg by (instance) (rate(node_cpu_seconds_total{mode="idle"}[5m])) * 100))',
                'memory_usage': 'avg((1 - (node_memory_MemAvailable_bytes / node_memory_MemTotal_bytes)) * 100)',
                'disk_usage': 'avg(100 - ((node_filesystem_avail_bytes{mountpoint="/"} / node_filesystem_size_bytes{mountpoint="/"}) * 100))'
            }
            
            summary = {}
            
            for metric_name, query in queries.items():
                try:
                    response = requests.get(
                        f"{self.base_url}/api/v1/query",
                        params={'query': query},
                        timeout=self.timeout
                    )
                    
                    if response.status_code == 200:
                        data = response.json()
                        result = data.get('data', {}).get('result', [])
                        
                        if result:
                            value = result[0].get('value', [None, '0'])[1]
                            summary[metric_name] = float(value) if value else 0
                        else:
                            summary[metric_name] = 0
                    else:
                        summary[metric_name] = 0
                        
                except Exception as e:
                    logger.warning(f"获取指标 {metric_name} 失败: {e}")
                    summary[metric_name] = 0
            
            logger.info("系统指标摘要获取完成")
            return summary
            
        except Exception as e:
            logger.error(f"获取系统指标异常: {e}")
            return {}
    
    def query_metric(self, query: str, time_range: str = '5m') -> List[Dict[str, Any]]:
        """执行PromQL查询"""
        try:
            params = {'query': query}
            
            # 如果需要时间范围查询
            if time_range:
                end_time = datetime.now()
                if time_range.endswith('m'):
                    minutes = int(time_range[:-1])
                    start_time = end_time - timedelta(minutes=minutes)
                elif time_range.endswith('h'):
                    hours = int(time_range[:-1])
                    start_time = end_time - timedelta(hours=hours)
                else:
                    start_time = end_time - timedelta(minutes=5)
                
                params.update({
                    'start': start_time.isoformat() + 'Z',
                    'end': end_time.isoformat() + 'Z',
                    'step': '30s'
                })
                
                endpoint = f"{self.base_url}/api/v1/query_range"
            else:
                endpoint = f"{self.base_url}/api/v1/query"
            
            response = requests.get(endpoint, params=params, timeout=self.timeout)
            
            if response.status_code == 200:
                data = response.json()
                return data.get('data', {}).get('result', [])
            else:
                logger.error(f"查询失败: HTTP {response.status_code}")
                return []
                
        except Exception as e:
            logger.error(f"查询异常: {e}")
            return []
    
    def get_prometheus_info(self) -> Dict[str, Any]:
        """获取Prometheus服务信息"""
        try:
            # 获取构建信息
            build_response = requests.get(f"{self.base_url}/api/v1/status/buildinfo", timeout=10)
            config_response = requests.get(f"{self.base_url}/api/v1/status/config", timeout=10)
            
            info = {
                'version': 'unknown',
                'status': 'unknown',
                'config_hash': 'unknown'
            }
            
            if build_response.status_code == 200:
                build_data = build_response.json()
                info['version'] = build_data.get('data', {}).get('version', 'unknown')
            
            if config_response.status_code == 200:
                config_data = config_response.json()
                info['config_hash'] = config_data.get('data', {}).get('yaml', '')[:50] + '...'
            
            info['status'] = 'connected'
            return info
            
        except Exception as e:
            logger.error(f"获取Prometheus信息失败: {e}")
            return {'version': 'unknown', 'status': 'disconnected', 'config_hash': 'unknown'}
    
    def _process_alert_data(self, alert: Dict[str, Any]) -> Dict[str, Any]:
        """处理告警数据标准化"""
        labels = alert.get('labels', {})
        annotations = alert.get('annotations', {})
        
        return {
            'alert_name': labels.get('alertname', '未知告警'),
            'instance': labels.get('instance', '未知实例'),
            'job': labels.get('job', '未知任务'),
            'severity': labels.get('severity', 'warning'),
            'summary': annotations.get('summary', ''),
            'description': annotations.get('description', ''),
            'started_at': alert.get('startsAt', ''),
            'state': alert.get('state', 'unknown'),
            'labels': labels,
            'annotations': annotations
        }